package com.zm.demo.javaagent.pre;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import lombok.extern.slf4j.Slf4j;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.ServiceLoader;

@Slf4j
public class PreAgent {

    static HashMap<String,String>enhanceMap = new HashMap<>();

    public static void premain(String agentArgs, Instrumentation inst){
        log.info("strat premain...");
        enhanceMap.put("className", "com.zm.demo.javaagent.test.UserController");
        enhanceMap.put("method", "health");
        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (className != null && enhanceMap.get("className").equalsIgnoreCase(className.replace("/","."))){
                    log.info("className = " + className);
                    try {
                        CtClass ctClass = ClassPool.getDefault().get(enhanceMap.get("className"));
                        CtMethod ctMethod = ctClass.getDeclaredMethod(enhanceMap.get("method"));
                        ctMethod.insertBefore("Long start  = System.nanoTime();");
                        ctMethod.insertAfter("System.out.println(\"cost: \" + System.nanoTime() - start + \"ms\");");
                        ctClass.detach();
                        return ctClass.toBytecode();
                    } catch (Exception e){

                    }
                }
                return new byte[0];
            }
        }, true);

        log.info("end premain...");
    }
}
